(* Yoann Padioleau
 *
 * Copyright (C) 2019-2024 Semgrep Inc.
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public License
 * version 2.1 as published by the Free Software Foundation, with the
 * special exception on linking described in file LICENSE.
 *
 * This library is distributed in the hope that it will be useful, but
 * WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the file
 * LICENSE for more details.
 *)
open Common
open Fpath_.Operators
module R = Rule
module PM = Core_match
module RP = Core_result
module E = Core_error
module OutJ = Semgrep_output_v1_t
module Log = Log_engine.Log

(*****************************************************************************)
(* Prelude *)
(*****************************************************************************)
(* Small wrapper around Match_search_mode and Match_tainting_mode *)

(*****************************************************************************)
(* Types *)
(*****************************************************************************)
(* This can be captured in Run_semgrep.ml *)
exception File_timeout of Rule_ID.t list

(* TODO make this one of the Semgrep_error_code exceptions *)
exception Multistep_rules_not_available

type timeout_config = {
  timeout : float;
  threshold : int;
  caps : < Cap.time_limit >;
  (* if eio is true, we are using a gc alarm based timeout mechanism. If
   * it's a None, we fall back to the signal-based one. *)
  eio : bool;
}

(*****************************************************************************)
(* Helpers *)
(*****************************************************************************)

let timeout_function (rule : Rule.t) (file : Fpath.t)
    (timeout : timeout_config option) f =
  let timeout, eio =
    match timeout with
    | None -> (None, false)
    | Some { eio; timeout; caps; threshold = _ } ->
        if timeout <= 0. then (None, eio) else (Some (timeout, caps), eio)
  in
  match
    Time_limit.set_timeout_opt ~name:"Match_rules.timeout_function" ~eio timeout
      f
  with
  | Some res -> Some res
  | None ->
      (* Note that we could timeout while testing the equality of two ASTs and
       * `busy_with_equal` will then erroneously have a `<> Not_busy` value. *)
      Log.err (fun m ->
          m "timeout for rule %s on file %s"
            (Rule_ID.to_string (fst rule.id))
            !!file);
      None

let group_relevant_rules rules (xconf : Match_env.xconfig) (xtarget : Xtarget.t)
    profiling =
  let interfile =
    (* TODO: Should be a field in 'xconf'. *)
    Option.is_some (Hook.get Pattern_vs_code.hook_find_possible_parents)
  in
  let non_sca_rules, sca_rules =
    rules
    |> List.partition_map (fun (r : Rule.t) ->
           match r.mode with
           | (`Extract _ | `Search _ | `Steps _ | `Taint _) as mode ->
               Left { r with mode }
           | `Join _ as mode ->
               Log.warn (fun m ->
                   m "Note: Join rule %a not avaliable via core" Rule_ID.pp
                     (fst r.id));
               Right { r with mode }
           | `SCA _ as mode ->
               Log.warn (fun m ->
                   m "Note: SCA rule %a not avaliable via core" Rule_ID.pp
                     (fst r.id));
               Right { r with mode })
  in
  let relevant, irrelevant, profiling =
    (* NOTE: Previously we updated xconf here based on the rule. I don't think
       `filter_irrelevant_rules` should be a rule-level option, so there
       shouldn't need to be any reason to update xconf here (since this is the
       only field we use here). Perhaps ideally rule-dependent options could
       have a different type (something like
       mk_rule_conf : Match_env.xconfig -> Rule.t -> rule_conf), options which
       ought to just be at the scan level can't have a dependency on rules. *)
    match xconf.filter_irrelevant_rules with
    (* NOTE: SCA rules are dropped. For details see discussion below in the
       filtered case. *)
    | NoPrefiltering -> (non_sca_rules, [], profiling)
    | CachedPrefilter f ->
        (* Rules we couldn't generate a prefilter for are always relevant *)
        let rules_with_prefilters, rules_with_no_prefilters =
          List.partition_map
            (fun r ->
              match f ~interfile (r :> Rule.t) with
              | Some prefilter -> Left (r, prefilter)
              | None -> Right r)
            non_sca_rules
        in
        let num_rules_with_prefilters = List.length rules_with_prefilters in
        Engine_metrics.Prefilter_metrics.record_rules_processed
          ~analyzer:xtarget.analyzer num_rules_with_prefilters;
        let profiling =
          Core_quick_profiling.add_rules_with_file_prefilters
            num_rules_with_prefilters profiling
        in
        let relevant_filtered, irrelevant_filtered =
          let%trace span = "prefilter.file" in
          let rules, prefilters = List_.split rules_with_prefilters in
          Log.info (fun m ->
              m "Performing rule prefiltering for %s"
                (Origin.to_string xtarget.path.origin));
          let prefilter_results =
            Tracing.with_span "prefilter.file.evaluation" ~__FILE__ ~__LINE__
              (fun _span ->
                Prefiltering.File.check_many prefilters
                  (Lazy_safe.force xtarget.lazy_content))
          in
          let rule_results = List_.combine_exn rules prefilter_results in
          let relevant_filtered, irrelevant_filtered =
            List.partition_map
              (fun ((rule, relevant) : _ Rule.rule_info * bool) ->
                if relevant then Left rule
                else (
                  Log.info (fun m ->
                      m "skipping rule %s for %s"
                        (Rule_ID.to_string (fst rule.id))
                        (Origin.to_string xtarget.path.origin));
                  Right rule))
              rule_results
          in
          (* Add span attributes for file prefiltering metrics *)
          Tracing.add_data_to_span span
            [
              ( "rules_with_prefilters_count",
                `Int (List.length rules_with_prefilters) );
              ("rules_filtered_count", `Int (List.length irrelevant_filtered));
              ("rules_kept_count", `Int (List.length relevant_filtered));
              ("analyzer", `String (Analyzer.to_string xtarget.analyzer));
              ("file_path", `String (Origin.to_string xtarget.path.origin));
            ];
          (relevant_filtered, irrelevant_filtered)
        in
        Engine_metrics.Prefilter_metrics.record_rules_skipped
          ~analyzer:xtarget.analyzer
          (List.length irrelevant_filtered);
        ( relevant_filtered @ rules_with_no_prefilters,
          irrelevant_filtered,
          profiling )
  in
  (* NOTE: SCA rules here are put into the irrelevant bucket. This is because
     we don't execute SCA rules in this codepath: SCA matching isn't available
     here in CE. If/when it is possible to execute them here, this will need to
     be updated. *)
  (`Relevant relevant, `Irrelevant (irrelevant @ sca_rules), profiling)

(* This function separates out rules into groups of taint rules by languages,
   all of the nontaint rules, and the rules which we skip due to prefiltering.
*)
let group_rules xconf rules xtarget profiling =
  let `Relevant relevant, `Irrelevant skipped_rules, profiling =
    group_relevant_rules rules xconf xtarget profiling
  in
  let relevant_taint_rules, relevant_nontaint_rules =
    List.partition_map
      (fun r ->
        match r.R.mode with
        | `Taint _ as mode -> Left { r with mode }
        | (`Extract _ | `Search _) as mode -> Right { r with mode }
        | `Steps _ ->
            Log.warn (fun m ->
                m "Step rule not handled: %s" (Rule.show_rule (r :> Rule.t)));
            raise Multistep_rules_not_available)
      relevant
  in
  (* Taint rules are only relevant to each other if they are meant to be
     analyzing the same language.
     So we group the taint rules by common language, before passing them
     to [Match_tainting_mode.check_rules].
  *)
  let relevant_taint_rules_groups =
    relevant_taint_rules
    |> List_.map (fun r -> (r.R.target_analyzer, r))
    |> Assoc.group_assoc_bykey_eff |> List_.map snd
  in
  let profiling =
    Core_quick_profiling.add_rules_selected
      (List.length relevant_taint_rules + List.length relevant_nontaint_rules)
      profiling
  in
  ( relevant_taint_rules_groups,
    relevant_nontaint_rules,
    skipped_rules,
    profiling )

(* Given a thunk [f] that computes the results of running the engine on a
 * single rule, this function simply instruments the computation on a single
 * rule with some boilerplate logic, like setting the last matched rule,
 * timing out if it takes too long, and producing a faulty match result in
 * that case.
 *
 * In particular, we need this to call [Match_tainting_mode.check_rules],
 * which will iterate over each rule in a different place, and so needs
 * access to this logic.
 *)
let per_rule_boilerplate_fn (timeout : timeout_config option) =
  let cnt_timeout = ref 0 in
  let rule_timeouts = ref [] in
  fun (file : Fpath.t) (rule : Rule.t) f ->
    let rule_id = fst rule.R.id in
    Domain.DLS.set Rule.last_matched_rule (Some rule_id);
    let res_opt =
      Profiling.measure
        (spf "real_rule:%s" (Rule_ID.to_string rule_id))
        (fun () ->
          (* here we handle the rule! *)
          timeout_function rule file timeout f)
    in
    match res_opt with
    | Some res -> res
    | None ->
        incr cnt_timeout;
        Stack_.push rule_id rule_timeouts;
        (match timeout with
        | Some { threshold; _ } when threshold > 0 && !cnt_timeout >= threshold
          ->
            raise (File_timeout !rule_timeouts)
        | _else_ -> ());
        let loc = Loc.first_loc_of_file file in
        let error = E.mk_error ~rule_id ~loc OutJ.Timeout in
        RP.mk_match_result []
          (E.ErrorSet.singleton error)
          (Core_profiling.empty_rule_profiling rule)

(*****************************************************************************)
(* Entry point *)
(*****************************************************************************)

let check ~matches_hook ~(timeout : timeout_config option)
    (xconf : Match_env.xconfig) (rules : Rule.rules) (xtarget : Xtarget.t) :
    Core_result.matches_single_file =
  let {
    path = { internal_path_to_content = file; _ };
    lazy_ast_and_errors;
    analyzer;
    _;
  } : Xtarget.t =
    xtarget
  in
  Log.info (fun m -> m "checking %s with %d rules" !!file (List.length rules));
  (match (!Profiling.profile, analyzer) with
  (* coupling: see Run_semgrep.xtarget_of_file() *)
  | Profiling.ProfAll, Analyzer.L (_lang, []) ->
      Log.debug (fun m ->
          m "forcing parsing of AST outside of rules, for better profile");
      Lazy_safe.force lazy_ast_and_errors |> ignore
  | _else_ -> ());

  let profiling =
    Core_quick_profiling.zero
    |> Core_quick_profiling.add_rules (List.length rules)
  in
  let per_rule_boilerplate_fn = per_rule_boilerplate_fn timeout file in

  (* We separate out the taint rules specifically, because we may want to
     do some rule-wide optimizations, which require analyzing more than
     just one rule at once.

     The taint rules are "grouped", see [group_rule] for more.

     TODO: use skipped_rules to call the commented skipped_target_of_rule?
  *)
  let taint_rules_groups, nontaint_rules, _skipped_rules, profiling =
    group_rules xconf rules xtarget profiling
  in
  let res_taint_rules, errors =
    taint_rules_groups
    |> List.fold_left
         (fun (res_acc, errors_acc) taint_rules ->
           let res, errors =
             Match_tainting_mode.check_rules ~matches_hook
               ~per_rule_boilerplate_fn taint_rules xconf xtarget
           in
           (List.rev_append res res_acc, errors_acc |> E.ErrorSet.union errors))
         ([], E.ErrorSet.empty)
  in
  let res_nontaint_rules =
    nontaint_rules
    |> List_.map (fun r ->
           let xconf =
             Match_env.adjust_xconfig_with_rule_options xconf r.R.options
           in
           per_rule_boilerplate_fn
             (r :> R.rule)
             (fun () ->
               Logs_.with_debug_trace ~__FUNCTION__
                 ~pp_input:(fun _ ->
                   "target: " ^ !!file ^ "\nruleid: "
                   ^ (r.id |> fst |> Rule_ID.to_string))
                 (fun () ->
                   (* dispatching *)
                   match r.R.mode with
                   | `Search _ as mode ->
                       Match_search_mode.check_rule { r with mode }
                         ~matches_hook xconf xtarget
                   | `Extract extract_spec ->
                       Match_search_mode.check_rule
                         { r with mode = `Search extract_spec.R.formula }
                         ~matches_hook xconf xtarget
                   | `Steps _ -> raise Multistep_rules_not_available)))
  in
  let res_total = res_taint_rules @ res_nontaint_rules in
  (* TODO: detect if a target was fully skipped because no rule
   * were irrelevant.
   * We used to report that for each rule independently via
   *
   * let skipped_target_of_rule (file_and_more : Xtarget.t) (rule : R.rule) :
   *     OutJ.skipped_target =
   *   let rule_id, _ = rule.id in
   *   let details =
   *     Some
   *       (spf
   *          "No need to perform matching because target does not contain \
   *           some elements necessary for the rule to match '%s'"
   *          (Rule_ID.to_string rule_id))
   *   in
   *   {
   *     path = file_and_more.path.internal_path_to_content;
   *     reason = Irrelevant_rule;
   *     details;
   *     rule_id = Some rule_id;
   *   }
   *
   * when skipped_target was part of profiling info, but now
   * skipped_target is only in the Core_result.t (and not in the
   * intermediate match_result).
   *)
  let profiling =
    Core_quick_profiling.add_rules_matched (List.length res_total) profiling
  in
  let res =
    RP.collate_rule_results xtarget.path.internal_path_to_content res_total
  in
  {
    res with
    quick_profiling = Some profiling;
    errors = res.errors |> E.ErrorSet.union errors;
  }
